package com.strp.mptcp.common;

import io.micrometer.core.instrument.*;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * MPTCP 性能指标收集器
 */
@Component
public class MPTCPMetrics {
    private final MeterRegistry registry;

    // 全局指标
    private final Counter totalBytesTransferred;
    private final Counter totalChunksTransferred;
    private final Counter totalRetransmissions;
    private final Gauge activeSubflows;
    private final Timer transferLatency;

    // 每个子流的指标
    private final Map<Integer, Counter> subflowBytesTransferred = new ConcurrentHashMap<>();
    private final Map<Integer, Counter> subflowRetransmissions = new ConcurrentHashMap<>();
    private final Map<Integer, Timer> subflowLatency = new ConcurrentHashMap<>();

    // 每个文件的传输进度
    private final Map<String, AtomicLong> fileProgress = new ConcurrentHashMap<>();

    public MPTCPMetrics(MeterRegistry registry) {
        this.registry = registry;

        // 初始化全局指标
        this.totalBytesTransferred = Counter.builder("mptcp.bytes.transferred.total")
                .description("Total bytes transferred")
                .register(registry);

        this.totalChunksTransferred = Counter.builder("mptcp.chunks.transferred.total")
                .description("Total chunks transferred")
                .register(registry);

        this.totalRetransmissions = Counter.builder("mptcp.retransmissions.total")
                .description("Total number of retransmissions")
                .register(registry);

        this.activeSubflows = Gauge.builder("mptcp.subflows.active",
                        () -> subflowBytesTransferred.size())
                .description("Number of active subflows")
                .register(registry);

        this.transferLatency = Timer.builder("mptcp.transfer.latency")
                .description("Transfer latency")
                .register(registry);
    }

    /**
     * 记录子流传输的字节数
     */
    public void recordBytesTransferred(int subflowId, long bytes) {
        totalBytesTransferred.increment(bytes);
        subflowBytesTransferred.computeIfAbsent(subflowId,
                        id -> Counter.builder("mptcp.bytes.transferred")
                                .tag("subflow", String.valueOf(id))
                                .register(registry))
                .increment(bytes);
    }

    /**
     * 记录数据块传输
     */
    public void recordChunkTransferred() {
        totalChunksTransferred.increment();
    }

    /**
     * 记录重传
     */
    public void recordRetransmission(int subflowId) {
        totalRetransmissions.increment();
        subflowRetransmissions.computeIfAbsent(subflowId,
                        id -> Counter.builder("mptcp.retransmissions")
                                .tag("subflow", String.valueOf(id))
                                .register(registry))
                .increment();
    }

    /**
     * 记录传输延迟
     */
    public void recordLatency(int subflowId, long milliseconds) {
        transferLatency.record(milliseconds, TimeUnit.MILLISECONDS);
        subflowLatency.computeIfAbsent(subflowId,
                        id -> Timer.builder("mptcp.latency")
                                .tag("subflow", String.valueOf(id))
                                .register(registry))
                .record(milliseconds, TimeUnit.MILLISECONDS);
    }

    /**
     * 更新文件传输进度
     */
    public void updateFileProgress(String fileId, long transferred, long total) {
        AtomicLong progress = fileProgress.computeIfAbsent(fileId,
                id -> new AtomicLong(0));
        progress.set(transferred);

        // 记录进度百分比
        double percentage = (double) transferred / total * 100;
        registry.gauge("mptcp.file.progress",
                Tags.of("file", fileId),
                percentage);
    }

    /**
     * 获取文件传输进度
     */
    public long getFileProgress(String fileId) {
        AtomicLong progress = fileProgress.get(fileId);
        return progress != null ? progress.get() : 0;
    }

    /**
     * 清理文件进度记录
     */
    public void clearFileProgress(String fileId) {
        fileProgress.remove(fileId);
        registry.remove(registry.find("mptcp.file.progress")
                .tag("file", fileId)
                .gauge());
    }

    /**
     * 获取子流传输统计
     */
    public SubflowStats getSubflowStats(int subflowId) {
        return new SubflowStats(
                subflowBytesTransferred.getOrDefault(subflowId,
                        Counter.builder("mptcp.bytes.transferred")
                                .tag("subflow", String.valueOf(subflowId))
                                .register(registry)).count(),
                subflowRetransmissions.getOrDefault(subflowId,
                        Counter.builder("mptcp.retransmissions")
                                .tag("subflow", String.valueOf(subflowId))
                                .register(registry)).count(),
                subflowLatency.getOrDefault(subflowId,
                        Timer.builder("mptcp.latency")
                                .tag("subflow", String.valueOf(subflowId))
                                .register(registry)).mean(TimeUnit.MILLISECONDS)
        );
    }

    /**
     * 子流统计信息
     */
    public record SubflowStats(
            double bytesTransferred,
            double retransmissions,
            double averageLatency
    ) {}
} 